package inference;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;

import com.analog.lyric.dimple.model.Discrete;

import model.ModelParameters;
import model.PredicateInfo;
import model.TableInformation;
import utils.Date;
import utils.IntPair;
import utils.Pair;
import utils.IntTriple;
import utils.TextUtils;

public class EntityFactor {

	
	public static Pair<int[][], double[]> createFactor(Discrete subject_variable, Discrete predicate_variable,
													   Discrete object_variable, ModelParameters params, int col,
													   HashMap<Integer, HashSet<Integer> > type_to_entity_idxs_map) {
		TableInformation table_info = params.getTableInformation();
		Object[] subjects = subject_variable.getDomain().getElements();
		Object[] predicates = predicate_variable.getDomain().getElements();
		int unk_pred_idx = predicates.length - 1;
		Object[] objects = object_variable.getDomain().getElements();
		int unk_object_idx = objects.length - 1;
		
		// Set up output data structures.
		int init_size = (int)(subjects.length * predicates.length * objects.length / 2.0);
		ArrayList<IntTriple> indices = new ArrayList<IntTriple>(init_size);
		ArrayList<Double> probabilities = new ArrayList<Double>(init_size);
		
		if (objects.length == 1) {
			for (int subj = 0; subj < subjects.length; subj++) {
				for (int pred = 0; pred < predicates.length; pred++) {
					indices.add(new IntTriple(subj, pred, 0));
					probabilities.add(1.0);
				}
			}
			return translateListsToPrimitives(indices, probabilities);
		}
		
		for (int subj = 0; subj < subjects.length; subj++) {
			indices.add(new IntTriple(subj, unk_pred_idx, unk_object_idx));
			probabilities.add(1.0);
		}
		
		PredicateInfo[] pred_infos = new PredicateInfo[predicates.length - 1];
		HashSet<Integer> valid_ranges = new HashSet<Integer>();
		for (int pred = 0; pred < predicates.length - 1; pred++) {
			int pred_id = ((Integer)(predicates[pred])).intValue();
			PredicateInfo pred_info = table_info.getPredicateInfo(pred_id); 
			pred_infos[pred] = pred_info;
			valid_ranges.add(pred_info.getRangeId());
		}
		
		
		
		for (int subj = 0; subj < subjects.length; subj++) {
			Integer subject_id = (Integer)(subjects[subj]);
//			String subject_name;
//			if (subj == unk_subject_idx) {
//				subject_name = "";
//			} else {
//				subject_name = table_info.getStringForId(subject_id);
//			} 
			HashSet<Integer> subj_types = table_info.getTypesForEntityId(subject_id);
//			
			HashMap<Integer, HashSet<Integer> > subj_fact_completions = table_info.getSubjFactCompletions(subject_id);
			
			for (int pred = 0; pred < predicates.length - 1; pred++) {
				Integer predicate_id = (Integer)(predicates[pred]);
				PredicateInfo pred_info = table_info.getPredicateInfo(predicate_id);
				if (!type_to_entity_idxs_map.containsKey(pred_info.getRangeId())) {
					indices.add(new IntTriple(subj, pred, unk_object_idx));
					probabilities.add(1.0);
					continue;
				}
				if ((subj_types != null && !subj_types.contains(pred_info.getDomainId()))) {
//					HashSet<Integer> type_entity_idxs = type_to_entity_idxs_map.get(pred_info.getRangeId());
					indices.add(new IntTriple(subj, pred, unk_object_idx));
					probabilities.add(1.0);
//					for (Integer entity_idx : type_entity_idxs) {
//						indices.add(new IntTriple(subj, pred, entity_idx));
//						probabilities.add((1.0) / (type_entity_idxs.size()));
//					}
////					indices.add(new IntTriple(subj, pred, unk_object_idx));
////					probabilities.add(1.0 / (type_entity_idxs.size() + 1));
////					probabilities.add(0.);
					continue;
				}
				
				boolean predicate_requires_raw_column = PredicateInfo.hasRawObjectColumn(table_info, pred_info);
				

				
				HashSet<Integer> object_completions;
				if (subj_fact_completions == null) {
					object_completions = null;
				} else {
					object_completions = subj_fact_completions.get(predicate_id);
				}
				
				HashSet<Integer> type_entity_idxs = type_to_entity_idxs_map.get(pred_info.getRangeId());
				if (objects.length == 1 || type_entity_idxs == null || type_entity_idxs.size() == 0) {
					indices.add(new IntTriple(subj, pred, 0));
					probabilities.add(1.0);
					continue;
				}
				double num_entities_of_type_for_cell = type_entity_idxs.size() + params.getApproximateTypeDomainSize();
				if (table_info.getStringForId(subject_id).equals("/m/03hc533") && table_info.getStringForId(predicate_id).equals("/business/employment_tenure/company")) {
					System.out.println("entities of type " + num_entities_of_type_for_cell);
				}
				Integer reverse_id = pred_info.getReverseId();
				PredicateInfo reverse_pred_info = null;
				if (reverse_id != null)
					reverse_pred_info = table_info.getPredicateInfo(reverse_id);
				
				HashMap<Integer, HashSet<Integer> > possible_reverse_subjs = new HashMap<Integer, HashSet<Integer> >(objects.length - 1);
				for (Integer type_entity_idx : type_entity_idxs) {
					possible_reverse_subjs.put(type_entity_idx,
							table_info.getSubjectsForObjAndPredPair((Integer)(objects[type_entity_idx.intValue()]), predicate_id));
				}
				if (pred_info.isFunctional()) {
					if (object_completions == null) {
						indices.add(new IntTriple(subj, pred, unk_object_idx));
						probabilities.add((double)(params.getApproximateTypeDomainSize()) / num_entities_of_type_for_cell);
						
						double uniform_probability = (1.0 / num_entities_of_type_for_cell);
						for (Integer entity_idx : type_entity_idxs) {
							indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
							probabilities.add(uniform_probability);
						}
					} else {  // Functional (subj, pred, obj) found in known facts.
						Integer true_object_id = object_completions.iterator().next();
						String true_object_string = table_info.getStringForId(true_object_id);
						Integer true_integer = TextUtils.parseIntegerString(true_object_string);
						String raw_cell_string = table_info.getStringForId((Integer)(objects[0]));
						
						int found_idx = indexOfEntity(objects, true_object_id);
						if (TextUtils.isMid(true_object_string) || found_idx != -1) {
							if (found_idx == -1) {
								indices.add(new IntTriple(subj, pred, unk_object_idx));
								probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
								
								double uniform_probability = params.getProbOfWrongFunctionalRel() / (type_entity_idxs.size());
								for (Integer entity_idx : type_entity_idxs) {
									indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
									probabilities.add(uniform_probability);
								}
							} else {
								if (objects.length == 2) {
									indices.add(new IntTriple(subj, pred, 0));
									probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
									indices.add(new IntTriple(subj, pred, 1));
									probabilities.add(params.getProbOfWrongFunctionalRel());
								} else {
									indices.add(new IntTriple(subj, pred, found_idx));
									probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
									
									num_entities_of_type_for_cell--;
									if (table_info.getStringForId(subject_id).equals("/m/03hc533") && table_info.getStringForId(predicate_id).equals("/business/employment_tenure/company")) {
										System.out.println("entities of type " + num_entities_of_type_for_cell);
									}
									
									indices.add(new IntTriple(subj, pred, unk_object_idx));
									probabilities.add(params.getProbOfWrongFunctionalRel() * 
													  params.getApproximateTypeDomainSize() / (double)num_entities_of_type_for_cell);
									
									double uniform_probability = params.getProbOfWrongFunctionalRel() / (num_entities_of_type_for_cell);
									for (Integer entity_idx : type_entity_idxs) {
										if (entity_idx == found_idx)
											continue;
										indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
										probabilities.add(uniform_probability);
									}
								}
							}
						} else if (raw_cell_string.equals(true_object_string)) {
							indices.add(new IntTriple(subj, pred, 0));
							probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
							indices.add(new IntTriple(subj, pred, 1));
							probabilities.add(params.getProbOfWrongFunctionalRel());
						} else if (TextUtils.isMid(raw_cell_string)) {
							indices.add(new IntTriple(subj, pred, unk_object_idx));
							probabilities.add(1.0);
						} else if (true_integer != null) {
							Integer cell_integer = TextUtils.parseIntegerString(raw_cell_string);
							if (cell_integer == null) {
								indices.add(new IntTriple(subj, pred, unk_object_idx));
								probabilities.add(1.0);
							} if (Math.abs(cell_integer - true_integer) < 0.001) {
								indices.add(new IntTriple(subj, pred, 0));
								probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
								indices.add(new IntTriple(subj, pred, 1));
								probabilities.add(params.getProbOfWrongFunctionalRel());
							} else {
								indices.add(new IntTriple(subj, pred, 0));
								probabilities.add(params.getProbOfWrongFunctionalRel());
								indices.add(new IntTriple(subj, pred, 1));
								probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
							}
						} else {
							Integer cell_integer = TextUtils.parseIntegerString(raw_cell_string);
							Date true_date = Date.parseDateString(true_object_string);
							if (cell_integer != null) {
								if (true_date == null) {
									indices.add(new IntTriple(subj, pred, 0));
									probabilities.add(params.getProbOfWrongFunctionalRel());
									indices.add(new IntTriple(subj, pred, 1));
									probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
								} else if (true_date.year == cell_integer.intValue()) {
									indices.add(new IntTriple(subj, pred, 0));
									probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
									indices.add(new IntTriple(subj, pred, 1));
									probabilities.add(params.getProbOfWrongFunctionalRel());
								} else {
									indices.add(new IntTriple(subj, pred, 0));
									probabilities.add(params.getProbOfWrongFunctionalRel());
									indices.add(new IntTriple(subj, pred, 1));
									probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
								}
							} else {
								Date cell_date = Date.parseTextDate(raw_cell_string);
								if (!(cell_date != null && true_date != null)) {
									indices.add(new IntTriple(subj, pred, 0));
									probabilities.add(params.getProbOfWrongFunctionalRel());
									indices.add(new IntTriple(subj, pred, 1));
									probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
								} else {
									if (cell_date.month == true_date.month && cell_date.year == true_date.year) {
										if (cell_date.day == 0 || true_date.day == 0 || cell_date.day == true_date.day) {
											indices.add(new IntTriple(subj, pred, 0));
											probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
											indices.add(new IntTriple(subj, pred, 1));
											probabilities.add(params.getProbOfWrongFunctionalRel());
										} else {
											indices.add(new IntTriple(subj, pred, 0));
											probabilities.add(params.getProbOfWrongFunctionalRel());
											indices.add(new IntTriple(subj, pred, 1));
											probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());	
										}
									} else {
										indices.add(new IntTriple(subj, pred, 0));
										probabilities.add(params.getProbOfWrongFunctionalRel());
										indices.add(new IntTriple(subj, pred, 1));
										probabilities.add(1.0 - params.getProbOfWrongFunctionalRel());
									}
								}
							}
						}
					}
				} else {  // Predicate is non functional.
					String raw_cell_string = table_info.getStringForId((Integer)(objects[0]));
					if (predicate_requires_raw_column && TextUtils.isMid(raw_cell_string)) {
						indices.add(new IntTriple(subj, pred, unk_object_idx));
						probabilities.add(1.0);
						continue;
					}
					int[] subj_is_correct = new int[type_entity_idxs.size()];
					for (int i = 0; i < type_entity_idxs.size(); i++) {
						subj_is_correct[i] = 0;
					}
					int num_incorrect = 0;
					int num_correct = 0;
					double total_cnt = 0;
					
					int ent = 0;
					for (Integer entity_idx : type_entity_idxs) {
							HashSet<Integer> possible_subjs = possible_reverse_subjs.get(entity_idx);
							
							if (possible_subjs == null || possible_subjs.size() == 0) {
								total_cnt += params.getLaplaceSmoothingAlpha();
								ent++;
								continue;
							} else {
								total_cnt += possible_subjs.size() + params.getLaplaceSmoothingAlpha();
							}
							if (reverse_pred_info != null && reverse_pred_info.isFunctional()) {
								if (possible_subjs.contains(subject_id)) {
									subj_is_correct[ent] = 1;
									total_cnt -= possible_subjs.size() + params.getLaplaceSmoothingAlpha();
									num_correct++;
								} else {
									subj_is_correct[ent] = -1;
									num_incorrect++;
									total_cnt -= possible_subjs.size() + params.getLaplaceSmoothingAlpha();
								}
							} else {
								if (object_completions != null && object_completions.contains((Integer)(objects[entity_idx.intValue()]))) {
									subj_is_correct[ent] = 1;
									total_cnt -= possible_subjs.size() + params.getLaplaceSmoothingAlpha();
									num_correct++;
								}
							}
						ent++;
					}
					total_cnt += (params.getApproximateTypeDomainSize() * params.getLaplaceSmoothingAlpha());
					
					if (num_incorrect != 0) {
						ent = 0;
						for (Integer entity_idx : type_entity_idxs) {
							if (subj_is_correct[ent] == -1) {
								indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
								probabilities.add(params.getProbOfWrongFunctionalRel() / num_incorrect);
							}
							ent++;
						}
					}
					if (table_info.getStringForId(subject_id).equals("/m/03hc533") && table_info.getStringForId(predicate_id).equals("/business/employment_tenure/company")) {
						System.out.println("reverse pred info " + reverse_pred_info);
					}
					
					if (predicate_requires_raw_column) {
							indices.add(new IntTriple(subj, pred, 0));
							probabilities.add(1.0 / num_entities_of_type_for_cell);
							indices.add(new IntTriple(subj, pred, 1));
							probabilities.add((double)(params.getApproximateTypeDomainSize()) / num_entities_of_type_for_cell);
					} else if (reverse_pred_info == null) {
						double numerator = 0.0;
						if (num_correct == 0) {
							numerator = 1.0;
						} else {
							numerator = params.getProbOfNewObjectInNonFunctionalRel();
						}
						indices.add(new IntTriple(subj, pred, unk_object_idx));
						probabilities.add(numerator * (double)(params.getApproximateTypeDomainSize() * params.getLaplaceSmoothingAlpha()) / total_cnt);
						
						double uniform_probability = (numerator / total_cnt);
						ent = 0;
						for (Integer entity_idx : type_entity_idxs) {
							if (subj_is_correct[ent] == 1) {
								indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
								probabilities.add((1.0 - params.getProbOfNewObjectInNonFunctionalRel()) / num_correct);
							} else if (subj_is_correct[ent] != -1) {
								indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
								probabilities.add((possible_reverse_subjs.get(entity_idx).size() + params.getLaplaceSmoothingAlpha()) * uniform_probability);
							}
							ent++;
						}
					} else {
						double numerator = 0.0;
						if (num_incorrect != 0) {
							if (num_correct == 0) {
								numerator = 1.0 - params.getProbOfWrongFunctionalRel();
							} else {
								numerator = params.getProbOfNewObjectInNonFunctionalRel();
							}
						} else if (num_correct == 0) {
							numerator = 1.0;
						} else {
							numerator = params.getProbOfNewObjectInNonFunctionalRel();
						}
						
						indices.add(new IntTriple(subj, pred, unk_object_idx));
						probabilities.add(numerator * (double)(params.getApproximateTypeDomainSize() * params.getLaplaceSmoothingAlpha()) / total_cnt);
						if (table_info.getStringForId(subject_id).equals("/m/03hc533") && table_info.getStringForId(predicate_id).equals("/business/employment_tenure/company")) {
							System.out.println("total cnt " + total_cnt);
						}
						double uniform_probability = (numerator / total_cnt);
						ent = 0;
						for (Integer entity_idx : type_entity_idxs) {
							if (subj_is_correct[ent] == 1) {
								indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
								if(num_incorrect == 0)
									probabilities.add((1.0 - params.getProbOfNewObjectInNonFunctionalRel()) / num_correct);
								else
									probabilities.add((1.0 - params.getProbOfNewObjectInNonFunctionalRel() - params.getProbOfWrongFunctionalRel()) / num_correct);
							} else if (subj_is_correct[ent] != -1) {
								indices.add(new IntTriple(subj, pred, entity_idx.intValue()));
								probabilities.add((possible_reverse_subjs.get(entity_idx).size() + params.getLaplaceSmoothingAlpha()) * uniform_probability);
							}
							ent++;
						}
					}
				}
			}
		}
		
		return translateListsToPrimitives(indices, probabilities);
	}
	
	
	private static Pair<int[][], double[]> translateListsToPrimitives(ArrayList<IntTriple> indices, ArrayList<Double> probabilities) {
		int[][] indices_arr = new int[indices.size()][3];
		double[] prob_arr = new double[probabilities.size()];
		
		Iterator<IntTriple> it = indices.iterator();
		Iterator<Double> it_prob = probabilities.iterator();
		for (int i = 0; i < indices.size(); i++) {
			IntTriple next = it.next();
			indices_arr[i] = next.toArray();
			prob_arr[i] = it_prob.next().doubleValue();
//			System.out.println(next.toString() + " " + prob_arr[i]);
		}
		
		return new Pair<int[][], double[]>(indices_arr, prob_arr);
	}
	
	private static int indexOfEntity(Object[] entities, Integer entity) {
		for (int i = 0; i < entities.length; i++) {
			if (((Integer)(entities[i])).equals(entity)) return i;
		}
		return -1;
	}
	
}
